import torch
from torch.utils.data import Dataset
import numpy as np
import cv2
import os

class ImageSimilarityDataset(Dataset):
    def __init__(self, image_dir, similarity_file):
        self.image_dir = image_dir
        self.similarities = np.load(similarity_file)
        self.image_files = sorted([f for f in os.listdir(image_dir) if f.endswith('.jpg')])
        
        # 归一化相似度值到0-1之间
        self.similarities = (self.similarities - np.min(self.similarities)) / (np.max(self.similarities) - np.min(self.similarities))
        
    def __len__(self):
        return len(self.similarities)

    def __getitem__(self, idx):
        img = cv2.imread(os.path.join(self.image_dir, self.image_files[idx]))
        similarity = self.similarities[idx]
        
        # 转换为RGB
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        
        # 转换为float32
        img = img.astype(np.float32) / 255.0
        
        # 转换为Tensor
        img = torch.from_numpy(img).permute(2, 0, 1)  # HWC to CHW
        similarity = torch.tensor(similarity, dtype=torch.float32)
        
        return img, similarity
